from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.messages import TextMessage
from autogen_agentchat.ui import Console
from autogen_core import CancellationToken
from autogen_ext.models.openai import OpenAIChatCompletionClient

# 测试 ollama 客户端
ollama_model_client = OpenAIChatCompletionClient(model="qwen2.5:7b", 
                                                 model_info={
                                                     "vision": False,
                                                     "function_calling": True,
                                                     "family": "Qwen3",
                                                     "structured_output": True,
                                                     "json_output": True,
                                                 },
                                                 api_key="ollama",
                                                 base_url="http://127.0.0.1:11434/v1")

# 测试 AssistantAgent
async def assistant_agent():
    agent = AssistantAgent(
        name="assistant",
        model_client=ollama_model_client,
        tools=[],
        system_message="Use tools to solve tasks.",
    )
    await Console(
        agent.on_messages_stream(
            [TextMessage(content="查询关于AutoGen的消息", source="user")],
            cancellation_token=CancellationToken(),
        )
    )


if __name__ == "__main__":
    import asyncio
    asyncio.run(assistant_agent())